# encoding: utf-8
# create train and val lmdb dataset individually
import os
import lmdb # install lmdb by "pip install lmdb"
import cv2
import numpy as np

dataPath = ['train', 'val']
baseRoot = 'F:/laibo/Data/CRNN_data/AllCRNNTestList/'
trainRoot = 'F:/laibo/Data/CRNN_data/AllCRNNTestList/annotation_train.txt'
valRoot = 'F:/laibo/Data/CRNN_data/AllCRNNTestList/annotation_val.txt'
lexiconRoot = 'F:/laibo/Data/CRNN_data/AllCRNNTestList/lexicon.txt'
outputPath = 'F:/laibo/Data/CRNN_data/lmdb/'


def checkImageIsValid(imageBin):
    if imageBin is None:
        return False
    imageBuf = np.fromstring(imageBin, dtype=np.uint8)
    # imageBuf = np.frombuffer(imageBin, dtype=np.uint8)
    img = cv2.imdecode(imageBuf, cv2.IMREAD_GRAYSCALE)
    imgH, imgW = img.shape[0], img.shape[1]
    if imgH * imgW == 0:
        return False
    return True


def writeCache(env, cache):
    with env.begin(write=True) as txn:
        for k, v in cache.items():
            # txn.put(k, v)
            # txn.put(k.encode(), str(v).encode())
            if isinstance(v, bytes):
                txn.put(k.encode(), v)
            else:
                txn.put(k.encode(), v.encode())


def createDataset(outputPath, imagePathList, labelList, lexiconList=None, checkValid=True):
    """
    Create LMDB dataset for CRNN training.
    ARGS:
        outputPath    : LMDB output path
        imagePathList : list of image path
        labelList     : list of corresponding groundtruth texts
        lexiconList   : (optional) list of lexicon lists
        checkValid    : if true, check the validity of every image
    """
    assert(len(imagePathList) == len(labelList))
    nSamples = len(imagePathList)
    # 注意map_size不要设置过大，byte为单位计算
    env = lmdb.open(outputPath, map_size=536870912)
    cache = {}
    cnt = 1
    for i in range(nSamples):
        imagePath = baseRoot + imagePathList[i]
        label = labelList[i]
        if not os.path.exists(imagePath):
            print('%s does not exist' % imagePath)
            continue
        with open(imagePath, 'rb') as f:
            imageBin = f.read()
        if checkValid:
            if not checkImageIsValid(imageBin):
                print('%s is not a valid image' % imagePath)
                continue

        imageKey = 'image-%09d' % cnt
        labelKey = 'label-%09d' % cnt
        cache[imageKey] = imageBin
        # cache[labelKey] = label
        # todo
        cache[labelKey] = lexiconList[int(label)]
        if lexiconList:
            lexiconKey = 'lexicon-%09d' % cnt
            # cache[lexiconKey] = ' '.join(lexiconList[i])
            cache[lexiconKey] = ' '.join(lexiconList[int(label)])
        if cnt % 1000 == 0:
            writeCache(env, cache)
            cache = {}
            print('Written %d / %d' % (cnt, nSamples))
        cnt += 1
    nSamples = cnt-1
    cache['num-samples'] = str(nSamples)
    writeCache(env, cache)
    print('Created dataset with %d samples' % nSamples)


if __name__ == '__main__':
    lexiconList = []
    f_lexicon = open(lexiconRoot, 'r', encoding='utf-8')
    for l in f_lexicon:
        lexiconList.append(l)
    for datapath in dataPath:
        if datapath == 'train':
            imagePathList = []
            labelList = []
            f_train = open(trainRoot, 'r')
            for l in f_train:
                imagepath, label = l.split(' ')
                imagePathList.append(imagepath)
                labelList.append(label)

            createDataset(outputPath + datapath, imagePathList, labelList, lexiconList)
        elif datapath == 'val':
            imagePathList = []
            labelList = []
            f_val = open(valRoot, 'r')
            for l in f_val:
                imagepath, label = l.split(' ')
                imagePathList.append(imagepath)
                labelList.append(label)

            createDataset(outputPath + datapath, imagePathList, labelList, lexiconList)
